from core.layer import Layer
from implement.layers.basic.batchnorm import BatchNorm
from implement.layers.conv.conv2d import Conv2d
from utils.functions_collect import relu


class BottleneckB(Layer):
    """A bottleneck layer that maintains the resolution of the feature map.
    Args:
        in_channels (int): Number of channels of input and output arrays.
        mid_channels (int): Number of channels of intermediate arrays.
    """

    def __init__(self, in_channels, mid_channels):
        super().__init__()

        self.conv1 = Conv2d(mid_channels, 1, 1, 0, nobias=True)
        self.bn1 = BatchNorm()
        self.conv2 = Conv2d(mid_channels, 3, 1, 1, nobias=True)
        self.bn2 = BatchNorm()
        self.conv3 = Conv2d(in_channels, 1, 1, 0, nobias=True)
        self.bn3 = BatchNorm()

    def forward(self, x):
        h = relu(self.bn1(self.conv1(x)))
        h = relu(self.bn2(self.conv2(h)))
        h = self.bn3(self.conv3(h))
        return relu(h + x)
